import sys
import re
import csv
import matplotlib.pyplot as plt

now_memory_size = 0
times = 0
x = []
y = []
y_max_static = []

dynamic_dict = {}
actor_dict = {}

max_dynamic_mem = 0
max_dynamic_graph = ""
max_static_mem = 0
peak_mem = 0


class User:
    def __init__(self):
        self.users = []

    def add_user(self, actor_name, index):
        self.users.append((actor_name, index))

    def __str__(self):
        dump = ""
        for user in self.users:
            dump += f"({user[0]}-index_{user[1]})"
        return dump


class Block:
    def __init__(self, start, end, size):
        self.start = start
        self.end = end
        self.size = size
        self.addr = ""
        self.type = ""
        self.actor = ""
        self.node = ""
        self.graph = ""
        self.user = User()

    def __str__(self):
        return f"start: {self.start}, end: {self.end}, size: {self.size}, addr: {self.addr}, type: {self.type}, " \
               f"actor: {self.actor}, node: {self.node}, graph: {self.graph}, user: {self.user}"


class Actor:
    def __init__(self, actor, graph, fm, stream):
        self.actor = actor
        self.graph = graph
        self.fm = fm
        self.stream = stream

    def __str__(self):
        return f"actor: {self.actor}, graph: {self.graph}, fm: {self.fm}, stream: {self.steam}"


class MemoryManagerInfo:
    def __init__(self, line_number, line):
        node = get_node(line)
        if node is None:
            raise ValueError(f"node not found, Line {line_number}: {line.strip()}")
        actor_name = get_actor_name(line)
        if actor_name is None:
            raise ValueError(f"actor name not found, Line {line_number}: {line.strip()}")
        graph = get_graph(line)
        if graph is None:
            raise ValueError(f"graph not found, Line {line_number}: {line.strip()}")
        device_address_class_ptr = get_device_address_class_ptr(line)
        if device_address_class_ptr is None:
            raise ValueError(f"device address class ptr not found, Line {line_number}: {line.strip()}")

        self.output_node = node
        self.actor_name = actor_name
        self.graph = graph
        self.device_address_class_ptr = device_address_class_ptr


block_list = []
actor_list = []
addr_block_dict = {}
manager_info_dict = {}
stream_dict = {}
fm_dict = {}


def get_actor_name(line):
    actor_name_pattern = r'actor name: ([A-Za-z0-9_]+)'
    actor_name_match = re.search(actor_name_pattern, line)
    if actor_name_match:
        actor_name = actor_name_match.group(1)
        if actor_name in actor_dict.keys():
            return actor_name + '_' + str(actor_dict[actor_name])
        else:
            raise ValueError(f"actor name {actor_name} not found in dict")
    else:
        return None


def get_size(line):
    size_pattern = r'size:\s*(\d+)'
    memory_size_match = re.search(size_pattern, line)
    if memory_size_match:
        number = memory_size_match.group(1)
        return number
    else:
        return -1


def get_graph(line):
    graph_pattern = r"graph: ([a-z0-9_]+)"
    graph_match = re.search(graph_pattern, line)
    if graph_match:
        graph = graph_match.group(1)
        return graph
    else:
        return None


def get_device_address_class_ptr(line):
    pattern = r"device address class ptr:\s*(\d+)"
    match = re.search(pattern, line)
    if match:
        address = match.group(1)
        return address
    else:
        return None


def get_node(line):
    node_pattern = r"node: (.*?),"
    node_match = re.search(node_pattern, line)
    if node_match:
        node = node_match.group(1)
        return node
    else:
        return None


def get_device_address(line):
    pattern = r"device address addr: (0x[0-9A-Fa-f]+)"
    match = re.search(pattern, line)
    if match:
        address = match.group(1)
        return address
    else:
        return None


def get_index(line):
    index_pattern = r"index: (\d+)"
    index_match = re.search(index_pattern, line)
    if index_match:
        index = index_match.group(1)
        return index
    else:
        return None


def memory_manager_info(line_number, line):
    if "Memory need alloc" in line:
        memory_manager_info = MemoryManagerInfo(line_number, line)
        manager_info_dict[memory_manager_info.device_address_class_ptr] = memory_manager_info


def memory_manager_actor(device_address_addr, line_number, line):
    device_address_class_ptr = get_device_address_class_ptr(line)
    if device_address_class_ptr is None:
        raise ValueError(f"device address class ptr not found, Line {line_number}: {line.strip()}")
    if device_address_class_ptr in manager_info_dict.keys():
        manager_info = manager_info_dict[device_address_class_ptr]
        actor_name = manager_info.actor_name
        graph = manager_info.graph
        node = manager_info.output_node

        if device_address_addr in addr_block_dict.keys():
            block = addr_block_dict[device_address_addr]
            block.actor = actor_name
            block.graph = graph
            block.node = node
        else:
            raise ValueError(f"device address addr {device_address_addr} not found in dict")
    else:
        raise ValueError(f"device address class ptr {device_address_class_ptr} not found in dict")


def memory_alloc_size(line_number, line):
    if "Memory pool alloc" in line:
        number = get_size(line)
        if number == -1:
            raise ValueError(f"memory alloc size not found, Line {line_number}: {line.strip()}")

        global now_memory_size, times
        block_list.append(Block(times, -1, number))

        now_memory_size = now_memory_size + int(number)
        y.append(now_memory_size)
        x.append(times)
        times = times + 1

        device_address_addr = get_device_address(line)
        if device_address_addr is None:
            raise ValueError(f"device address addr not found, Line {line_number}: {line.strip()}")

        addr_block_dict[device_address_addr] = block_list[-1]
        block_list[-1].addr = device_address_addr

        global max_static_mem, peak_mem
        total_memory_pattern = r'total mem:\s*(\d+)'
        total_memory_match = re.search(total_memory_pattern, line)
        if total_memory_match:
            total_memory = total_memory_match.group(1)
            total_memory = int(total_memory)
            if total_memory > max_static_mem:
                max_static_mem = total_memory
            y_max_static.append(max_static_mem)
        else:
            raise ValueError(f"total mem not found, Line {line_number}: {line.strip()}")

        peak_memory_pattern = r'peak mem:\s*(\d+)'
        peak_memory_match = re.search(peak_memory_pattern, line)
        if peak_memory_match:
            peak_memory = peak_memory_match.group(1)
            peak_memory = int(peak_memory)
            if peak_memory > peak_mem:
                peak_mem = peak_memory
        else:
            raise ValueError(f"peak mem not found, Line {line_number}: {line.strip()}")


def memory_free_size(line_number, line):
    if "Memory pool free" in line:
        number = get_size(line)
        if number == -1:
            raise ValueError(f"memory free size not found, Line {line_number}: {line.strip()}")
        global now_memory_size, times
        now_memory_size = now_memory_size - int(number)
        y.append(now_memory_size)
        y_max_static.append(max_static_mem)
        x.append(times)

        device_address_addr = get_device_address(line)
        if device_address_addr is None:
            raise ValueError(f"device address addr not found, Line {line_number}: {line.strip()}")
        if device_address_addr in addr_block_dict.keys():
            addr_block_dict[device_address_addr].end = times
        else:
            raise ValueError(
                f"device address addr {device_address_addr} not found in dict, Line {line_number}: {line.strip()}")
        times = times + 1


def actor_name(line_number, line):
    if "launch actor name" in line:
        actor_name_pattern = r'launch actor name: ([A-Za-z0-9_]+)'
        actor_name_match = re.search(actor_name_pattern, line)
        if actor_name_match:
            actor_name = actor_name_match.group(1)
            if actor_name not in actor_dict.keys():
                actor_dict[actor_name] = 0
            else:
                actor_dict[actor_name] = actor_dict[actor_name] + 1
        else:
            raise ValueError(f"actor name not found, Line {line_number}: {line.strip()}")


def feature_memory_size(line_number, line):
    if "feature memory" in line:
        feature_memory_pattern = r'feature memory:\s*(\d+)'
        feature_memory_match = re.search(feature_memory_pattern, line)
        graph = get_graph(line)
        if graph is None:
            raise ValueError(f"graph not found, Line {line_number}: {line.strip()}")
        if feature_memory_match:
            number = feature_memory_match.group(1)
            dynamic_dict[times] = int(number)
            fm_dict[graph] = int(number)

            global max_dynamic_mem, max_dynamic_graph
            if int(number) > max_dynamic_mem:
                max_dynamic_mem = int(number)
                max_dynamic_graph = graph
        else:
            raise ValueError(f"feature memory size not found, Line {line_number}: {line.strip()}")


def memory_type(line_number, line):
    if "alloc type" in line:
        memory_type_pattern = r"alloc type: ([A-Za-z]+)"
        memory_type_match = re.search(memory_type_pattern, line)
        if memory_type_match:
            memory_type = memory_type_match.group(1)
            device_address_addr = get_device_address(line)
            if device_address_addr is None:
                raise ValueError(f"device address addr not found, Line {line_number}: {line.strip()}")
            if device_address_addr in addr_block_dict.keys():
                addr_block_dict[device_address_addr].type = memory_type
            else:
                raise ValueError(
                    f"device address addr {device_address_addr} not found in dict, Line {line_number}: {line.strip()}")

            if memory_type == "MemoryManagerActor":
                memory_manager_actor(device_address_addr, line_number, line)

            graph = get_graph(line)
            if graph is not None:
                addr_block_dict[device_address_addr].graph = graph

            node = get_node(line)
            if node is not None:
                addr_block_dict[device_address_addr].node = node

            actor_name = get_actor_name(line)
            if actor_name is not None:
                addr_block_dict[device_address_addr].actor = actor_name

        else:
            raise ValueError(f"memory type not found, Line {line_number}: {line.strip()}")


def memory_use(line_number, line):
    if "Memory use" in line:
        actor_name = get_actor_name(line)
        if actor_name is None:
            raise ValueError(f"actor name not found, Line {line_number}: {line.strip()}")
        device_address_addr = get_device_address(line)
        if device_address_addr is None:
            raise ValueError(f"device address addr not found, Line {line_number}: {line.strip()}")
        index = get_index(line)
        if index is None:
            raise ValueError(f"index not found, Line {line_number}: {line.strip()}")

        if device_address_addr in addr_block_dict.keys():
            addr_block_dict[device_address_addr].user.add_user(actor_name, index)


def graph_stream(line_number, line):
    if "stream" in line:
        graph = get_graph(line)
        if graph is None:
            raise ValueError(f"graph not found, Line {line_number}: {line.strip()}")
        stream_pattern = r'stream:\s*(\d+)'
        stream_match = re.search(stream_pattern, line)
        if stream_match:
            number = stream_match.group(1)
            stream_dict[graph] = int(number)
        else:
            raise ValueError(f"feature memory size not found, Line {line_number}: {line.strip()}")


def end_launch(line_number, line):
    if "end launch" in line:
        graph = get_graph(line)
        if graph is None:
            raise ValueError(f"graph not found, Line {line_number}: {line.strip()}")
        actor_name = get_actor_name(line)
        if actor_name is None:
            raise ValueError(f"actor name not found, Line {line_number}: {line.strip()}")
        stream = -1
        fm = -1
        if graph in stream_dict.keys():
            stream = stream_dict[graph]
        else:
            raise ValueError(f"graph {graph} not found in stream dict, Line {line_number}: {line.strip()}")
        if graph in fm_dict.keys():
            fm = fm_dict[graph]
        else:
            raise ValueError(f"graph {graph} not found in fm dict, Line {line_number}: {line.strip()}")

        actor_list.append(Actor(actor_name, graph, fm, stream))


def create_block_csv():
    csv_data = [['Start', 'End', 'Size', 'Addr', 'MemType', 'ActorName', 'NodeName', 'GraphName', 'User']]
    for block in block_list:
        csv_data.append([block.start, block.end, block.size, block.addr, block.type, block.actor, block.node,
                         block.graph, block.user])

    csv_filename = 'blocks.csv'
    with open(csv_filename, mode='w', newline='') as file:
        writer = csv.writer(file)
        writer.writerows(csv_data)


def create_actor_csv():
    all_stream = 0
    stream_set = set()
    csv_data = [['ActorName', 'Graph', 'FeatureMemory', 'Stream', 'AllStream']]
    for actor in actor_list:
        if actor.graph not in stream_set:
            stream_set.add(actor.graph)
            all_stream += actor.stream
        csv_data.append([actor.actor, actor.graph, actor.fm, actor.stream, all_stream])

    csv_filename = 'actor.csv'
    with open(csv_filename, mode='w', newline='') as file:
        writer = csv.writer(file)
        writer.writerows(csv_data)


def create_memory_chart():
    plt.clf()
    plt.figure(figsize=(40, 10))
    plt.plot(x, y, marker='o', linestyle='-', color='b', label='Static Memory')
    plt.plot(x, y_max_static, marker='o', linestyle='-', color='g', label='Max Pool Memory')

    dyn_x = []
    dyn_y = []
    for i in range(len(y)):
        if i in dynamic_dict.keys():
            dyn_x.append(i)
            dyn_y.append(dynamic_dict[i])

    plt.scatter(dyn_x, dyn_y, color='r', marker='+', label='Dynamic Memory')

    plt.title('Memory')
    plt.xlabel('x')
    plt.ylabel('Byte')
    plt.legend()
    plt.savefig('memory.png')


def analyze_file(filename):
    try:
        with open(filename, 'r') as file:
            lines = file.readlines()
            for line_number, line in enumerate(lines, start=1):
                memory_alloc_size(line_number, line)
                memory_free_size(line_number, line)
                feature_memory_size(line_number, line)
                memory_type(line_number, line)
                actor_name(line_number, line)
                memory_manager_info(line_number, line)
                memory_use(line_number, line)
                graph_stream(line_number, line)
                end_launch(line_number, line)
    except FileNotFoundError:
        print(f"Error: File '{filename}' not found.")

    create_memory_chart()
    create_block_csv()
    create_actor_csv()

    print(f"max dynamic mem: {max_dynamic_mem} bytes, graph: {max_dynamic_graph}")
    print(
        f"max static mem: {max_static_mem} bytes, peak used: {peak_mem} bytes, memory pool utilization: {peak_mem * 1.0 / max_static_mem}")


if __name__ == "__main__":
    if len(sys.argv) != 2:
        print("Usage: python analysis.py filename")
        sys.exit(1)

    file_name = sys.argv[1]
    analyze_file(file_name)
